package week6

import (
	"sync"
	"time"
)

type Counter struct {
	BucketSize        int
	BucketNumber      int
	CounterTime       int64
	Buckets           []map[interface{}]int64
	BucketWindowStart []int64
	Peak              int
	Mutex             *sync.RWMutex
}

func NewCounter(bucketSize int, bucketNumber int) *Counter {
	c := &Counter{
		BucketSize:        bucketSize,
		BucketNumber:      bucketNumber,
		CounterTime:       int64(bucketSize * bucketNumber),
		Buckets:           make([]map[interface{}]int64, bucketNumber),
		BucketWindowStart: make([]int64, bucketNumber),
		Peak:              0,
		Mutex:             &sync.RWMutex{},
	}
	return c
}

// Increment the counter in the current bucket by one for the given key.
func (c *Counter) Increment(k interface{}) {
	c.Mutex.Lock()
	defer c.Mutex.Unlock()
	c.getCurrentBucket()[k]++
}

// Reset force a reset of all rolling counters (clear all buckets) so that statistics start being gathered from scratch.
func (c *Counter) Reset() {
	c.Mutex.Lock()
	defer c.Mutex.Unlock()
	c.Buckets = make([]map[interface{}]int64, c.BucketNumber)
	c.BucketWindowStart = make([]int64, c.BucketNumber)
	c.Peak = 0
}

// GetRollingSum get the sum of all buckets in the rolling counter for the given key.
func (c *Counter) GetRollingSum(k interface{}) int64 {
	var sum int64 = 0

	c.Mutex.Lock()
	defer c.Mutex.Unlock()
	c.getCurrentBucket()
	for _, v := range c.Buckets {
		if v == nil {
			continue
		}
		sum += v[k]
	}
	return sum
}

// GetValueOfLatestBucket get the value of the latest (current) bucket in the rolling counter for the given key.
func (c *Counter) GetValueOfLatestBucket(k interface{}) int64 {
	c.Mutex.Lock()
	defer c.Mutex.Unlock()
	return c.getCurrentBucket()[k]
}

func (c *Counter) getCurrentBucket() map[interface{}]int64 {
	now := time.Now().UnixNano()

	if c.Buckets[c.Peak] != nil && (now < c.BucketWindowStart[c.Peak]+int64(c.BucketSize)) {
		return c.Buckets[c.Peak]
	}

	peak := -1
	var latest int64 = 0
	for i := len(c.Buckets) - 1; i > 0; i-- {
		if now-(c.BucketWindowStart[i]+c.CounterTime) > 0 {
			c.Buckets[i] = nil
		}
		if latest < c.BucketWindowStart[i] {
			peak = i
			latest = c.BucketWindowStart[i]
		}
	}

	newPeak := (peak + 1) % c.BucketNumber
	newWindowStart := latest + int64(c.BucketSize)
	for {
		newPeak = (peak + 1) % c.BucketNumber
		newWindowStart += int64(c.BucketSize)
		c.BucketWindowStart[newPeak] = newWindowStart
		c.Buckets[newPeak] = make(map[interface{}]int64)

		if now < (newWindowStart + int64(c.BucketSize)) {
			c.Peak = newPeak
			break
		}
	}

	return c.Buckets[c.Peak]
}
